﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using System.Windows.Forms;

namespace bio2WF
{
    class NJTreeBuilder
    {
        public double scale { get; set; }
        private RichTextBox rtb;

        public class TreeBranch
        {
            public List<TreeBranch> branches { get; set; }
            public string name { get; set; }
            public double edgeLength { get; set; }
            public Boolean isLeaf { get; set; }

            public TreeBranch(List<TreeBranch> branches, string name, double edgeLength, Boolean isLeaf)
            {
                this.branches = branches;
                this.name = name;
                this.edgeLength = edgeLength;
                this.isLeaf = isLeaf;
            }
        }

        Dictionary<string, Dictionary<string, double>> distances;
        List<TreeBranch> treeBranches;

        public NJTreeBuilder(Dictionary<string, Dictionary<string, double>> distances)
        {
            this.distances = distances;
            treeBranches = getInitialStarTree(distances.Keys);
            scale = 1;
        }

        public List<TreeBranch> getInitialStarTree(Dictionary<string, Dictionary<string, double>>.KeyCollection leafs)
        {
            List<TreeBranch> resultList = new List<TreeBranch>();
            foreach (string leaf in leafs)
            {
                resultList.Add(new TreeBranch(new List<TreeBranch>(), leaf, 1 * scale, true));
            }
            return resultList;
        }

        public string generateTreeInNewickFormat(RichTextBox rtb)
        {
            this.rtb = rtb;
            return buildTree(distances);
        }

        public string buildTree(Dictionary<string, Dictionary<string, double>> distances)
        {
            printMatrix(distances);
            int i = 1;
            while (treeBranches.Count - 2 != 0)
            {
                // wylicz macierz Q wedlug wzoru
                Dictionary<string, Dictionary<string, double>> qMatrix = getQMatrix(distances);
                rtb.Text += ("Q-Matrix:");
                printMatrix(qMatrix);

                // wybierz pare z najmniejsza wartoscia w macierzy Q
                string[] smallestPair = getSmallestPair(qMatrix);
                rtb.Text += ("\nNajmniejsze wartosci dla: " + smallestPair[0] + " i " + smallestPair[1]);

                // dolacz wezel z najmniejsza para i obliczonymi wartosciami
                TreeBranch smallestBranch1 = getTreeBranchByName(treeBranches, smallestPair[0]);
                TreeBranch smallestBranch2 = getTreeBranchByName(treeBranches, smallestPair[1]);

                TreeBranch newBranch = new TreeBranch(new List<TreeBranch>(), "U" + (i++).ToString(), 1 * scale, false);
                double distanceFromSmallest1ToEveryOther = getDistanceToEveryOther(smallestPair[0], distances) / (treeBranches.Count - 2);
                double distanceFromSmallest2ToEveryOther = getDistanceToEveryOther(smallestPair[1], distances) / (treeBranches.Count - 2);
                smallestBranch1.edgeLength = ((distances[smallestPair[0]][smallestPair[1]]) / 2 + (distanceFromSmallest1ToEveryOther - distanceFromSmallest2ToEveryOther) / 2) * scale;
                smallestBranch2.edgeLength = ((distances[smallestPair[0]][smallestPair[1]]) / 2 + (distanceFromSmallest2ToEveryOther - distanceFromSmallest1ToEveryOther) / 2) * scale;
                treeBranches.Remove(smallestBranch1);
                treeBranches.Remove(smallestBranch2);
                newBranch.branches.Add(smallestBranch1);
                newBranch.branches.Add(smallestBranch2);
                treeBranches.Add(newBranch);
                rtb.Text += ("\n\nAktualne drzewo:\n(" + convertToNewickFormat(treeBranches) + ");");

                // uaktualnij macierz dystansu
                Dictionary<string, double> newTaxon = new Dictionary<string, double>();
                double smallestPairDistance = distances[smallestPair[0]][smallestPair[1]];
                distances.Remove(smallestPair[0]);
                distances.Remove(smallestPair[1]);
                foreach (string taxon in distances.Keys)
                {
                    double length = distances[taxon][smallestPair[0]] + distances[taxon][smallestPair[1]] - smallestPairDistance;
                    newTaxon.Add(taxon, length / 2);
                    distances[taxon].Remove(smallestPair[0]);
                    distances[taxon].Remove(smallestPair[1]);
                    distances[taxon].Add(newBranch.name, length / 2);
                }
                distances.Add(newBranch.name, newTaxon);
                distances[newBranch.name].Add(newBranch.name, 0);

                rtb.Text += ("\n\nNowa macierz dystansu:\n");
                printMatrix(distances);
                rtb.Text += ("\n\n=================================\n\n");
            }

            // polacz odstatnie dwa wezly;
            if (treeBranches.ElementAt(0).isLeaf)
            {
                treeBranches.Reverse();
            }
            treeBranches.ElementAt(1).edgeLength = distances[treeBranches.ElementAt(0).name][treeBranches.ElementAt(1).name] * scale;
            treeBranches.ElementAt(0).branches.Add(treeBranches.ElementAt(1));
            treeBranches.Remove(treeBranches.ElementAt(1));
            treeBranches.ElementAt(0).edgeLength = 0;

            rtb.Text += ("\n\nKoncowe drzewo:\n\n(" + convertToNewickFormat(treeBranches) + ");");

            return "(" + convertToNewickFormat(treeBranches) + ");";
        }

        public Dictionary<string, Dictionary<string, double>> getQMatrix(Dictionary<string, Dictionary<string, double>> distances)
        {
            Dictionary<string, Dictionary<string, double>> qMatrix = new Dictionary<string, Dictionary<string, double>>();
            foreach (string taxon in distances.Keys)
            {
                qMatrix.Add(taxon, new Dictionary<string, double>());
                foreach (string taxon2 in distances.Keys)
                {
                    int r = distances.Keys.Count;
                    double sumOfDistancesFromTaxonToEveryOther = 0;
                    double sumOfDistancesFromTaxon2ToEveryOther = 0;
                    foreach (string taxon3 in distances.Keys)
                    {
                        sumOfDistancesFromTaxonToEveryOther += distances[taxon][taxon3];
                        sumOfDistancesFromTaxon2ToEveryOther += distances[taxon2][taxon3];
                    }
                    double value = distances[taxon][taxon2] - sumOfDistancesFromTaxonToEveryOther / (r - 2) - sumOfDistancesFromTaxon2ToEveryOther / (r - 2);
                    qMatrix[taxon].Add(taxon2, value);
                }
            }
            foreach (string taxon in distances.Keys)
            {
                qMatrix[taxon][taxon] = 0;
            }
            return qMatrix;
        }


        public void printMatrix(Dictionary<string, Dictionary<string, double>> matrix)
        {
            foreach (string taxon in distances.Keys)
            {
                rtb.Text += ("\t" + taxon);
            }
            rtb.Text += ("\n");
            foreach (string taxon in distances.Keys)
            {
                rtb.Text += (taxon);
                foreach (string taxon2 in distances.Keys)
                {
                    rtb.Text += ("\t" + String.Format("{0:F2}", matrix[taxon][taxon2]));
                }
                rtb.Text += ("\n");
            }
        }

        public string[] getSmallestPair(Dictionary<string, Dictionary<string, double>> qMatrix)
        {
            double smallestValue = double.MaxValue;
            string[] smallestPair = { "", "" };
            foreach (string taxon in distances.Keys)
            {
                foreach (string taxon2 in distances.Keys)
                {
                    if (qMatrix[taxon][taxon2] < smallestValue)
                    {
                        smallestValue = qMatrix[taxon][taxon2];
                        smallestPair[0] = taxon;
                        smallestPair[1] = taxon2;
                    }
                }
            }
            return smallestPair;
        }

        public TreeBranch getTreeBranchByName(List<TreeBranch> treeBranches, string name)
        {
            foreach (TreeBranch branch in treeBranches)
            {
                if (branch.name == name)
                {
                    return branch;
                }
            }
            return null;
        }

        public double getDistanceToEveryOther(string taxonName, Dictionary<string, Dictionary<string, double>> distances)
        {
            double result = 0;
            foreach (string taxon in distances.Keys)
            {
                result += distances[taxonName][taxon];
            }
            return result;
        }

        public string convertToNewickFormat(List<TreeBranch> treeBranches)
        {
            string result = "";
            int i = 0;
            foreach (TreeBranch treeBranch in treeBranches)
            {
                if (treeBranch.isLeaf)
                {
                    result += treeBranch.name + ":" + String.Format("{0:F2}", treeBranch.edgeLength).Replace(',', '.');
                }
                else
                {
                    if (treeBranch.edgeLength != 0) result += "(";
                    result += convertToNewickFormat(treeBranch.branches);
                    if (treeBranch.edgeLength != 0) result += "):" + String.Format("{0:F2}", treeBranch.edgeLength).Replace(',', '.');
                }
                if (i++ != treeBranches.Count - 1)
                {
                    result += ",";
                }
            }
            return result;
        }
    }
}
